from enum import Enum
import torch
from semilearn.core.hooks import Hook


class MultiheadAgreementTypes(Enum):
    LowConfDisagreement = 1
    HighConfDisagreement = 2
    DiffConfDisagreement = 3
    LowConfAgreement = 4
    HighConfAgreement = 5
    DiffConfAgreement = 6
    NA = 7


def check_apm_agreement(value):
    if value in [MultiheadAgreementTypes.HighConfAgreement]:
        return 1 # agreement with valid pseudo-label
    elif value in [MultiheadAgreementTypes.DiffConfDisagreement, MultiheadAgreementTypes.DiffConfAgreement]:
        return 0 # disagreement with valid pseudo-label
    elif value in [MultiheadAgreementTypes.LowConfDisagreement, MultiheadAgreementTypes.HighConfDisagreement, MultiheadAgreementTypes.LowConfAgreement, MultiheadAgreementTypes.NA]:
        return -1 # no valid pseudo-label
    else:
        raise ValueError(f'{value} is not a valid MultiheadAgreementTypes.')


class MultiMatchThresholdingHook(Hook):
    def __init__(self, args) -> None:
        super().__init__()
        self.args = args
        self.num_heads = args.num_heads
        self.num_classes = args.num_classes
        self.ulb_dest_len = args.ulb_dest_len
        self.smoothness = args.smoothness
        self.apm_percentile = args.apm_percentile

        self.apm = torch.zeros((self.num_heads, self.ulb_dest_len, self.num_classes)).cuda(args.gpu)

        # apm_cutoff[head_id][class_id] = APM threshold for class_id when head_id is excluded
        self.apm_cutoff = torch.zeros((self.num_heads, self.num_classes)).cuda(args.gpu)

        # predictions for weakly augmented data for each head (head pseudolabels)
        self.weak_labels = torch.ones((self.num_heads, self.ulb_dest_len)).cuda(args.gpu) * -1

        # weak labels that exceed APM threshold for agreement
        self.weak_agreement_mask = torch.zeros((self.num_heads, self.ulb_dest_len), dtype=torch.bool).cuda(args.gpu)

        # predictions for strongly augmented data for each head
        self.strong_labels = torch.ones((self.num_heads, self.ulb_dest_len), dtype=torch.int64).cuda(args.gpu) * -1

        # labels generated by APM values, where heads didn't reach an agreement
        self.apm_labels = torch.ones((self.num_heads, self.ulb_dest_len), dtype=torch.int).cuda(args.gpu) * -1

        # MultiheadAgreementTypes
        self.agreement_types = torch.zeros((self.num_heads, self.ulb_dest_len), dtype=torch.int).cuda(args.gpu)

    def update(self, algorithm, logits_x_ulb_w, logits_x_ulb_s, idx_ulb, head_id):
        epoch = algorithm.epoch
        
        max_probs, pseudo_labels = torch.topk(logits_x_ulb_w, 2)
        for i, idx in enumerate(idx_ulb.tolist()):
            pm = logits_x_ulb_w[i] - max_probs[i][0]
            pm[pseudo_labels[i][0]] = max_probs[i][0] - max_probs[i][1]

            self.apm[head_id][idx] = pm * self.smoothness / (1 + epoch) + self.apm[head_id][idx] * (1 - self.smoothness / (1 + epoch))

            self.weak_labels[head_id][idx] = pseudo_labels[i][0]

        self.strong_labels[head_id][idx_ulb] = torch.max(logits_x_ulb_s, dim=-1)[1]

    def get_apm_label(self, algorithm, idx, label1, label2, head_id, head_id1, head_id2):
        self.weak_agreement_mask[head_id][idx] = False

        # We don't have default apm_cutoff values for the first epoch
        if algorithm.epoch == 0:
            _multihead_agreement_type = MultiheadAgreementTypes.NA
            self.agreement_types[head_id][idx] = _multihead_agreement_type.value
            return -1, _multihead_agreement_type.value, check_apm_agreement(_multihead_agreement_type)
    
        apm_pass1 = self.apm[head_id1][idx][label1] >= self.apm_cutoff[head_id][label1]
        apm_pass2 = self.apm[head_id2][idx][label2] >= self.apm_cutoff[head_id][label2]

        if label1 == label2:
            # At least one head APM needs to exceed threshold
            if apm_pass1 and apm_pass2:
                self.weak_agreement_mask[head_id][idx] = True
                _multihead_agreement_type = MultiheadAgreementTypes.HighConfAgreement
                self.agreement_types[head_id][idx] = _multihead_agreement_type.value
                return label1, _multihead_agreement_type.value, check_apm_agreement(_multihead_agreement_type)
            
            elif apm_pass1 or apm_pass2:
                self.weak_agreement_mask[head_id][idx] = True
                _multihead_agreement_type = MultiheadAgreementTypes.DiffConfAgreement
                self.agreement_types[head_id][idx] = _multihead_agreement_type.value
                return label1, _multihead_agreement_type.value, check_apm_agreement(_multihead_agreement_type)
            
            else:
                _multihead_agreement_type = MultiheadAgreementTypes.LowConfAgreement
                self.agreement_types[head_id][idx] = _multihead_agreement_type.value
                return -1, _multihead_agreement_type.value, check_apm_agreement(_multihead_agreement_type)

        
        # Both labels exceed the thresholds => We don't know what to choose, so ignore
        if apm_pass1 and apm_pass2:
            _multihead_agreement_type = MultiheadAgreementTypes.HighConfDisagreement
            self.agreement_types[head_id][idx] = _multihead_agreement_type.value
            return -1, _multihead_agreement_type.value, check_apm_agreement(_multihead_agreement_type)
        
        if apm_pass1:
            _multihead_agreement_type = MultiheadAgreementTypes.DiffConfDisagreement
            self.agreement_types[head_id][idx] = _multihead_agreement_type.value
            self.apm_labels[head_id][idx] = label1
            return label1, _multihead_agreement_type.value, check_apm_agreement(_multihead_agreement_type)
        
        if apm_pass2:
            _multihead_agreement_type = MultiheadAgreementTypes.DiffConfDisagreement
            self.agreement_types[head_id][idx] = _multihead_agreement_type.value
            self.apm_labels[head_id][idx] = label2
            return label2, _multihead_agreement_type.value, check_apm_agreement(_multihead_agreement_type)
        
        # Neither of the 2 labels has passed the threshold
        _multihead_agreement_type = MultiheadAgreementTypes.LowConfDisagreement
        self.agreement_types[head_id][idx] = _multihead_agreement_type.value
        return -1, _multihead_agreement_type.value, check_apm_agreement(_multihead_agreement_type)

    def after_train_step(self, algorithm):
        if (self.args.num_recalibrate_iter > 0 and self.every_n_iters(algorithm, self.args.num_recalibrate_iter)) \
            or (self.args.num_recalibrate_iter == 0 and self.end_of_epoch(algorithm)):

            algorithm.print_fn("Recomputing apm_cutoff...")

            for head_id, head_id1, head_id2 in [(0, 1, 2), (1, 0, 2), (2, 0, 1)]:
                for label in range(self.num_classes):
                    # Both heads need to choose the selected label
                    mask = self.weak_labels[head_id1] == self.weak_labels[head_id2]
                    mask = mask & (self.weak_labels[head_id1] == label)

                    # Nothing was found => Use 0 as apm_cutoff
                    if 1 not in mask:
                        self.apm_cutoff[head_id][label] = 0
                        break
                    
                    # Merge apm values for head1 and head2 for the current label
                    apm_set1 = self.apm[head_id1][mask]
                    apm_set2 = self.apm[head_id2][mask]
                    apm_set = torch.cat((apm_set1[:, label], apm_set2[:, label]), dim=0)

                    # Use apm_percentile as the new threshold
                    if self.args.no_low:
                        self.apm_cutoff[head_id][label] = torch.quantile(apm_set.float(), self.apm_percentile)
                    else:
                        self.apm_cutoff[head_id][label] = max(0, torch.quantile(apm_set.float(), self.apm_percentile))

            # reset all APM labels to -1
            self.apm_labels = torch.ones((self.num_heads, self.ulb_dest_len), dtype=torch.int).cuda(self.args.gpu) * -1
            self.agreement_types = torch.zeros((self.num_heads, self.ulb_dest_len), dtype=torch.int).cuda(self.args.gpu)
